use std::{
    collections::BTreeMap,
    fmt::{Debug, Display},
};

#[cfg(test)]
extern crate quickcheck;

use ordered_float::NotNan;

use tokio::sync::mpsc::UnboundedSender;

use crate::mec::Entropic;

use crate::{
    DiscreteDistribution, argmax,
    common::{
        DataOfInterest, DecodeError, NextTokenDistributionProvider, compute_marginals,
        log_mu_table, make_uniform, pick_istar,
    },
    min_entropy_joint_distribution_sparse, order_descending,
};

pub async fn decode<A: Clone + Eq + Display, W: Clone + Eq + PartialEq, V>(
    ss: &mut (dyn Iterator<Item = W> + Sync + Send),
    //message_space_size: usize,
    message_space_labels: &[A],

    cnext: &mut impl NextTokenDistributionProvider<W, V>,

    // this has to be known!
    message_parts: usize,

    cnext_start: V,

    logger: Option<&slog::Logger>,
    tx: Option<&UnboundedSender<DataOfInterest>>,
) -> Result<Vec<A>, anyhow::Error>
where
    A: Debug,
    W: Debug,
{
    cnext.reset(cnext_start).await?;

    let message_space_size = message_space_labels.len();

    let should_stop = |mu: &[DiscreteDistribution<_>]| {
        mu.iter()
            .all(|distro: &DiscreteDistribution<A>| distro.d.entropy() < 0.1)
    };

    let mut mu = (0..message_parts)
        .map(|_| make_uniform(message_space_size, Some(message_space_labels.to_vec())))
        .collect::<Vec<DiscreteDistribution<A>>>();
    let mut done = false;
    let mut output = vec![];

    let mut itercount = 0;

    if let Some(tx) = tx {
        tx.send(DataOfInterest::DecodeStart())?;
    };

    log_mu_table(tx, &mu);

    while !done {
        if let Some(logger) = logger {
            info!(logger, "--------- decode iteration {}", itercount);
        }
        done = decode_step(
            &mut mu,
            cnext,
            ss,
            message_space_size,
            should_stop,
            logger,
            tx,
        )
        .await?;

        itercount += 1;
    }

    // if let Some(logger) = logger {
    //     info!(logger, "decoded mu: {}", mu);
    // }

    for distr in mu {
        output.push(argmax(&distr))
    }

    if let Some(logger) = logger {
        info!(logger, "result = {:#?}", output);
    }

    Ok(output)
}

async fn decode_step<A: Clone + Display + Debug, W: Clone + Eq + Debug, V>(
    mu: &mut [DiscreteDistribution<A>],
    cnext: &mut impl NextTokenDistributionProvider<W, V>,
    ss: &mut (dyn Iterator<Item = W> + Sync + Send),
    message_space_size: usize,
    should_stop: impl Fn(&[DiscreteDistribution<A>]) -> bool,
    logger: Option<&slog::Logger>,
    tx: Option<&UnboundedSender<DataOfInterest>>,
) -> Result<bool, anyhow::Error> {
    if let Some(logger) = logger {
        let mu_entropy_values_by_i = mu.iter().map(|x| x.d.entropy()).enumerate();
        mu_entropy_values_by_i.for_each(|(i, h)| {
            debug!(logger, "distro mu_{} has entropy {:?}", i, h);
        });
    }

    let istar = pick_istar(mu);

    if let Some(logger) = logger {
        info!(logger, "picking istar={}", istar);
    }

    tx.map(|tx| tx.send(DataOfInterest::IStar(istar)));

    let mut next_token_distribution = cnext.get().await?;

    order_descending(&mut next_token_distribution); // this should always be OK!

    // let n = next_token_distribution.d.len(); // for logging only

    let sj: W = ss.next().ok_or(DecodeError::NoNextCodeSymbol())?;

    // if let Some(logger) = logger {
    //     info!(
    //         logger,
    //         "n={}, next token distro C(C_j | ...) = {:?}", n, next_token_distribution
    //     );
    //     info!(logger, "mu distro = {:?}", mu[istar]);
    // }

    let (gammaj, swap_p_q) =
        min_entropy_joint_distribution_sparse(next_token_distribution.d, mu[istar].d.clone());

    if let Some(logger) = logger {
        info!(logger, "gammaj: {:?}, swap={}", gammaj, swap_p_q);
        info!(logger, "{:?}", compute_marginals(&gammaj));
    }

    // println!("{:?} {:?}", next_token_distribution.labels, sj);

    let j_sj = next_token_distribution
        .labels
        .as_ref()
        .unwrap()
        .iter()
        .position(|label| *label == sj)
        .ok_or(DecodeError::NoSuchToken(format!(
            "{:?} {:?}",
            next_token_distribution.labels, sj
        )))?;

    cnext.select(sj.clone()).await?;

    let weights_map_column = BTreeMap::from_iter(
        gammaj
            .iter()
            .filter(|(_, (j, jprime))| *(if swap_p_q { jprime } else { j }) == j_sj)
            .map(|(x, (iprime, i))| (*(if swap_p_q { iprime } else { i }), *x)),
    );

    let weights_column_with_indices =
        (0..message_space_size).map(|i| (weights_map_column.get(&i).unwrap_or(&0.0), i));

    let updated_mu = DiscreteDistribution {
        d: weights_column_with_indices
            .map(|x| *x.0)
            .collect::<Vec<_>>(),
        labels: mu[istar].labels.clone(),
    };

    // Normalize to get the conditional probability
    let sum: f64 = updated_mu.d.iter().sum();
    let mut updated_mu = DiscreteDistribution {
        d: updated_mu.d.iter().map(|x| x / sum).collect::<Vec<_>>(),
        labels: updated_mu.labels,
    };

    if let Some(logger) = logger {
        info!(
            logger,
            "entropy of updated mu[{}] = {}",
            istar,
            updated_mu.d.entropy()
        );
    }

    // Sort by probabilities in non-increasing order again
    let mut indices = (0..updated_mu.d.len()).collect::<Vec<_>>();
    indices.sort_by_key(|&i| std::cmp::Reverse(NotNan::new(updated_mu.d[i]).unwrap()));

    if let Some(logger) = logger {
        debug!(logger, "indices sorted by key = {:?}", indices);
    }

    updated_mu.labels = updated_mu
        .labels
        .map(|v| indices.iter().map(|&i| v[i].clone()).collect::<Vec<_>>());

    updated_mu.d = indices.iter().map(|&i| updated_mu.d[i]).collect();

    // if let Some(logger) = logger {
    //     debug!(
    //         logger,
    //         "the updated mu is now {:?} (labels {:?})", updated_mu.d, updated_mu.labels
    //     );
    // }

    mu[istar] = updated_mu;

    log_mu_table(tx, mu);

    Ok(should_stop(mu))
}
